#region Credits
//
// Credit for original work:
// http://www.codeproject.com/Articles/12072/C-NET-DNS-query-component
// Bdev.Net.Dns by Rob Philpott, Big Developments Ltd rob@bigdevelopments.co.uk
// 
#endregion

using System;
using System.Net;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;

namespace Windows.Dns
{
    /// <summary>
    /// Summary description for Dns.
    /// </summary>
    public static class Resolver
    {
        private const int DnsPort = 53;
        private static int uniqueId;


        /// <summary>
        /// The principal look up function, which sends a request message to the given
        /// DNS server and collects a response. This implementation re-sends the message
        /// via UDP up to two times in the event of no response/packet loss
        /// </summary>
        /// <param name="request">The logical request to send to the server</param>
        /// <param name="dnsServer">The IP address of the DNS server we are querying</param>
        /// <param name="timeoutInSeconds">The amount of time to wait for a response (in seconds)</param>
        /// <returns>The logical response from the DNS server or null if no response</returns>
        public static async Task<Response> Lookup(Request request, IPAddress dnsServer, int timeoutInSeconds = 2)
        {
            // check the inputs
            if (request == null) throw new ArgumentNullException("request");
            if (dnsServer == null) throw new ArgumentNullException("dnsServer");

            return await Task.Factory.StartNew(() =>
                                                   {
                                                       // We will not catch exceptions here, rather just refer them to the caller

                                                       // create an end point to communicate with
                                                       var server = new IPEndPoint(dnsServer, DnsPort);


                                                       SocketRequest sr = null;
                                                       var mre = new ManualResetEvent(false);
                                                       DataReceived = bts =>
                                                                          {
                                                                              // ReSharper disable AccessToModifiedClosure
                                                                              if (sr == bts)
                                                                                  // ReSharper restore AccessToModifiedClosure
                                                                              {
                                                                                  mre.Set();
                                                                              }
                                                                          };

                                                       // send the request and get the response
                                                       //byte[] responseMessage =;
                                                       sr = UdpTransfer(server, request);
                                                       mre.WaitOne(timeoutInSeconds*1000);
                                                       return sr.Response;
                                                   });
        }

        private static SocketRequest UdpTransfer(EndPoint server, Request requestMessage)
        {
            var sr = new SocketRequest {Request = requestMessage};

            // we'll be send and receiving a UDP packet
            var socket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);

            var socketEventArg = new SocketAsyncEventArgs();
            socketEventArg.Completed += SocketEventArgCompleted;
            socketEventArg.RemoteEndPoint = server;
            socketEventArg.UserToken = sr;

            sr.Socket = socket;
            socket.ConnectAsync(socketEventArg);
            return sr;

        }

        // A single callback is used for all socket operations. 
        // This method forwards execution on to the correct handler 
        // based on the type of completed operation
        private static void SocketEventArgCompleted(object sender, SocketAsyncEventArgs e)
        {
            switch (e.LastOperation)
            {
                case SocketAsyncOperation.Connect:
                    ProcessConnect(e);
                    break;
                case SocketAsyncOperation.Receive:
                    ProcessReceive(e);
                    break;
                case SocketAsyncOperation.Send:
                    ProcessSend(e);
                    break;
                default:
                    throw new Exception("Invalid operation completed");
            }
        }

        // Called when a ConnectAsync operation completes
        private static void ProcessConnect(SocketAsyncEventArgs e)
        {
            if (e.SocketError == SocketError.Success)
            {
                // Successfully connected to the server
                var sock = e.UserToken as SocketRequest;
                if (sock == null) return;

                byte[] requestMessage;
                try
                {

                    // get the message
                    requestMessage = sock.Request.GetMessage();

                    // firstly, uniquely mark this request with an id
                    unchecked
                    {
                        // substitute in an id unique to this lookup, the request has no idea about this
                        requestMessage[0] = (byte) (uniqueId >> 8);
                        requestMessage[1] = (byte) uniqueId;
                    }

                }
                finally
                {
                    // increase the unique id
                    uniqueId++;
                }


                // Send 
                e.SetBuffer(requestMessage, 0, requestMessage.Length);

                var willRaiseEvent = sock.Socket.SendAsync(e);
                if (!willRaiseEvent)
                {
                    ProcessSend(e);
                }
            }
            else
            {
                throw new SocketException((int) e.SocketError);
            }
        }

        private static event Action<SocketRequest> DataReceived;

        // Called when a ReceiveAsync operation completes
        // </summary>
        private static void ProcessReceive(SocketAsyncEventArgs e)
        {
            if (e.SocketError == SocketError.Success)
            {
                var sock = e.UserToken as SocketRequest;
                if (sock == null) return;

                // Received data from server
                var data = e.Buffer;
                var resp = new Response(data);
                sock.Response = resp;
                DataReceived(sock);

                // Data has now been sent and received from the server. 
                // Disconnect from the server
                sock.Socket.Shutdown(SocketShutdown.Send);
                sock.Socket.Close();
            }
            else
            {
                throw new SocketException((int) e.SocketError);
            }
        }


        // Called when a SendAsync operation completes
        private static void ProcessSend(SocketAsyncEventArgs e)
        {
            if (e.SocketError == SocketError.Success)
            {
                var sock = e.UserToken as SocketRequest;
                if (sock == null) return;

                //Read data sent from the server
                e.SetBuffer(new byte[512], 0, 512);
                bool willRaiseEvent = sock.Socket.ReceiveAsync(e);
                if (!willRaiseEvent)
                {
                    ProcessReceive(e);
                }
            }
            else
            {
                throw new SocketException((int) e.SocketError);
            }
        }

        private class SocketRequest
        {
            public Socket Socket { get; set; }
            public Request Request { get; set; }
            public Response Response { get; set; }
        }
    }

}
